# coding=utf-8

from net import cls_net
from preprocess import get_data
import config 
import torch
import torch.nn as nn
import torch.optim as optim
import time
import os
import logging

def train():
    logging.info("Start training.")
    # get data
    trainloader = get_data(config.data_path)
    # get HRNet
    net = cls_net()
    # define a optimizer
    optimizer = optim.Adam(net.parameters(), lr=config.lr)
    # load model if exists
    if os.path.exists(config.model_load_path) is not True:
        criterion = nn.CrossEntropyLoss()
    else:
        checkpoint = torch.load(path)
        net.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        initepoch = checkpoint['epoch']
        criterion = checkpoint['loss']

    # train
    total = correct = 0
    for epoch in range(config.epoch): 
        timestart = time.time()

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            optimizer.zero_grad()
            # forward, backward and step
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 200 == 199:
                logging.info('[%d, %5d] loss: %.3f' %
                    (epoch + 1, i + 1, running_loss / 200))
                running_loss = 0.0
                # get accuaracy of train dataset
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                logging.info('Accuracy of the network on the %d train images: %.3f %%' % (total,
                        100.0 * correct / total))


                # save the model
                torch.save({'epoch':epoch,
                            'model_state_dict':net.state_dict(),
                            'optimizer_state_dict':optimizer.state_dict(),
                            'loss':criterion
                            }, config.model_save_path)

        logging.info('epoch %d cost %3f sec' %(epoch,time.time()-timestart))

    logging.info('Finished Training')

if __name__ == "__main__":
    logging.basicConfig(filename="./log/train.log", filemode="w", format="%(asctime)s %(name)s:%(levelname)s:%(message)s", datefmt="%D-%M-%Y %H:%M:%S", level=logging.DEBUG)
    train()
